function [optima, RegionSet, SampleSet, RegionSampleId] = PRS_MMO_v1( Problem, AlgorithmM )
% PRS_MMO: Partition-Based Random Search for Multimodal Optimization for minimization problems

%INPUT
%   Problem.
%       Dimension: scalar, the dimension of the decision variables
%       Domain: 1 -by- ~ vector, features of the domain
%       Partition: function handle of partitioning strategy
%           [RegionNum, Region, RegionSampleId] = Partition(Region, SampleSet)
%               RegionNum: scalar, number of new partitioned regions
%               Region = [partition depth, partitionable, features ...]
%                   partition depth = 0 for the whole domain
%                   partitionable: scalar, 1: partitionable, 0: non-partitionable
%               RegionSampleId: ~-by-1 cell, the sample id that belongs to each region
%               SampleSet: [SampleID, objective value, x1, x2, ...]
%       Sampling: function handle of sampling in a particular region
%           [Sample] = Sampling(n, region)
%               n: number of samples
%               region: 1 -by- ~ vector, [features...]
%               Sample: n -by- (1+d) matrix, [objective value, x1, x2, ...]
%   AlgorithmM.
%       n0: scalar, first stage sampling size
%       QuantileLevel: scalar, <0.5, the quantile level of promisinhg region definition
%       NewBudget: scalar, added budget size at each stage
%       MaxSampleSize: scalar, maximal sample size in one region
%       Radius: scalar, the radius distance that be regarded as neighbor
%       StopCriteria: [X,Y]
%           [1,Y]: reach maximum budget Y
%           [2,Y]: reach the required number of optima Y
%           [3,Y]: the number of optima is not changed in Y iterations
%OUTPUT
%   optima: ~ -by- (2+d) matrix, the final optimal solution set,
%         [sample id, objective value, x1, x2, ...]
%   RegionSet: ~ -by- (2+~) matrix, the whole subregion set
%         [partition depth, partitionable, features ...]
%   SampleSet: ~ -by- (1+d) matrix, the original sample set
%         [objective value, x1, x2, ...]
%   RegionSampleId: ~ -by- ~ cell, the id of the samples in each region

%% Initialization
% PROBLEM PARAMETERS
d = Problem.Dimension;
domain = [0, 1, Problem.Domain]; % [partition depth, partitionable, features ...]
Partition = Problem.Partition;
Sampling = Problem.Sampling;
% ALGORITHM PARAMETERS
n0 = AlgorithmM.n0;
alpha = AlgorithmM.QuantileLevel;
Delta = AlgorithmM.NewBudget;
MaxSampleSize = AlgorithmM.MaxSampleSize;
StopCriteria = AlgorithmM.StopCriteria;
N = (StopCriteria(1)==1)*StopCriteria(2) + (StopCriteria(1)~=1)*1000000; % the maximal budget size
radius = AlgorithmM.radius;

%% Optimization
% PRE-PROCESSING
RegionSet = zeros(round(N/n0), length(domain));
RegionSet(1,:) = domain;
RegionsNum = 1; % the number of exist regions
CurrentMaxPartition = 0;  % the maximum depth of partitioning at current iteration
GroupMean = zeros(round(N/n0), 1); % group sample mean
GroupStd = zeros(round(N/n0), 1); % group sample std
GroupN1 = zeros(round(N/n0), 1); % real group sample size
GroupN1_mod = zeros(round(N/n0), 1); % modified group sample size
GroupWeight = zeros(round(N/n0), 2); % [the weight of each group, the group is changed or not];
RegionSampleId = cell(round(N/n0), 1); % id of group samples
SampleSet = zeros(N, 2+d);  % [SampleID,objective value,x1,x2,...]
SampleSet(:,1) = 1:size(SampleSet,1);
OptCandidate=zeros(N,1);  % 1: the possible optima, 0: not possible
N0 = 0; % budget allocated

% START OPTIMIZATION
NumOpt = 0;  % the current optimal value
NumIterationRepeat = 1; % the number of iterations that the number of optima is not changed
PartitionList = zeros(Delta,1);
PartitionList(1) = 1;
PartitionListNum = 1;
% subsequent iterations
while N0 < N...
    && (StopCriteria(1)~=2 || NumOpt<StopCriteria(2))...
    && (StopCriteria(1)~=3 || NumOpt==0 || NumIterationRepeat<StopCriteria(2))...
    && sum(RegionSet(1:RegionsNum,2))>0
%% partition partitionable regions
    i = 1;
    while i <= PartitionListNum
        iPartitionRegion = PartitionList(i);
        [NewRegionNum, NewRegion, NewRegionSampleId]...
            = Partition(RegionSet(iPartitionRegion,:), SampleSet(RegionSampleId{iPartitionRegion},:));
        NewRegionId = [iPartitionRegion,(RegionsNum+1):(RegionsNum+NewRegionNum-1)]; % id of new-partitioned region
        RegionSet(NewRegionId,:) = NewRegion;
        RegionSampleId(NewRegionId) = NewRegionSampleId;
        RegionsNum = RegionsNum+NewRegionNum-1; % update total region number
        % update partition state
        PartitionDepth_temp = max(NewRegion(:,1));
        if PartitionDepth_temp>CurrentMaxPartition
            CurrentMaxPartition = PartitionDepth_temp;
        end
        % update the optima set
        nonpartitionable_temp = find(NewRegion(:,2)==0);
        for j=1:length(nonpartitionable_temp)
            OptCandidate(RegionSampleId{NewRegionId(nonpartitionable_temp(j))})=1;
        end
        if StopCriteria(1)~=1 && ~isempty(nonpartitionable_temp)
            optima = extract;
            NumOpt_old = NumOpt;
            NumOpt = size(optima,1);
            if NumOpt_old ~= NumOpt
                NumIterationRepeat = 1;
            end
        end
        % first stage sampling and group information updating
        for j = NewRegionId
            GroupN1(j) = length(RegionSampleId{j});
            GroupWeight(j,2) = 1;
            if GroupN1(j)>=MaxSampleSize && RegionSet(j,2)==1
                PartitionListNum = PartitionListNum+1;
                PartitionList(PartitionListNum) = j;
            else
                % first stage sampling
                if RegionSet(j,2)==0 % non-partitionable
                    Add_n = min(1-GroupN1(j),N-N0);
                else
                    Add_n = min(n0-GroupN1(j),N-N0);
                end
                if Add_n>0
                    NewSampleId = ((N0+1):(N0+Add_n))';
                    SampleSet(NewSampleId,2:end) = Sampling(Add_n,RegionSet(j,3:end));
                    N0 = N0+Add_n;
                    RegionSampleId{j} = [RegionSampleId{j};NewSampleId];
                    GroupN1(j) = GroupN1(j)+Add_n;
                end
                % update group information
                Samplevalue_temp = SampleSet(RegionSampleId{j},2);
                GroupMean(j) = mean(Samplevalue_temp);
                GroupStd(j) = sqrt(var(Samplevalue_temp));
            end
        end
        i = i+1;
    end
    PartitionListNum = 0;
    
%% modify the BAQM equations
    ModConst = RegionSet(1:RegionsNum,1)./CurrentMaxPartition;
    GroupN1_mod(1:RegionsNum)=max(2,round(GroupN1(1:RegionsNum).*ModConst));
%% additional budget
    Add_n = min(N-N0,Delta);
    n2 = BAQM(Add_n,alpha,GroupMean(1:RegionsNum)',GroupStd(1:RegionsNum)',GroupN1_mod(1:RegionsNum)', RegionSet(1:RegionsNum,2)');
    RegionId_temp = find(n2>0);
    for i = RegionId_temp
        NewSampleId = ((N0+1):(N0+n2(i)))';
        SampleSet(NewSampleId,2:end) = Sampling(n2(i),RegionSet(i,3:end));
        N0 = N0+n2(i);
        RegionSampleId{i} = [RegionSampleId{i};NewSampleId];
        mean_temp = GroupMean(i);
        GroupMean(i) = (GroupMean(i)*GroupN1(i)+sum(SampleSet(NewSampleId,2)))/(GroupN1(i)+n2(i));
        GroupStd(i) = sqrt((GroupStd(i)^2*(GroupN1(i)-1)...
            +(GroupMean(i)-mean_temp)^2*GroupN1(i)...
            +sum((SampleSet(NewSampleId,2)-GroupMean(i)).^2))/(GroupN1(i)+n2(i)-1));
        GroupN1(i) = GroupN1(i)+n2(i);
        if GroupN1(i)>=MaxSampleSize && RegionSet(i,2)==1
            PartitionListNum = PartitionListNum+1;
            PartitionList(PartitionListNum) = i;
        end
        GroupWeight(i,2) = 1;
    end
    NumIterationRepeat = NumIterationRepeat+1;
end
%% final results
if StopCriteria(1)==1
    [optima] = extract();
%     optima = [];
end
RegionSet((RegionsNum+1):end, :) = [];
SampleSet((N0+1):end, :) = [];
SampleSet(:,1)=[];
RegionSampleId((RegionsNum+1):end, :) = [];

%% Plots for 2 dimension box constraint problems
% if d==2
%     figure()
%     hold on
%     for i=1:RegionsNum
%         rectangle('Position',[RegionSet(i,[3,5]),RegionSet(i,[4,6])-RegionSet(i,[3,5])])
%     end
%     scatter(SampleSet(:,2),SampleSet(:,3),2,SampleSet(:,1),'filled');
%     hold off
% end


%% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% FUNCTIONS
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    function [ n2 ] = BAQM( N_new, Percentage, mu_hat, sigma_hat, n1, partitionable )
    %BAQM Budget Allocation for Quantile Minimization
    %INPUT
    %   N_new: the new simulation budget
    %   Percentage: the percentage of good points of interest
    %   mu_hat: 1-by-K vector, sample mean of each group
    %   sigma_hat: 1-by-K vector, sample standard deviation of each group
    %   n1: 1-by-K vector, number of observations in each group at stage one
    %   partitionable: 1-by-K vector, 1: the region is partitionable
    %OUTPUT
    %   n2: 1-by-K vector, additional budgets need to be allocated to each group
        zero_tolerance = 1e-16;
        y0 = mu_hat+sigma_hat.*norminv(Percentage,0,1);
        [y,b] = min(y0); % threshold and the best group
        NumRegion = size(mu_hat,2);
        if GroupWeight(b,2)==1
            GroupWeight(1:NumRegion,2) = 1; % all regions should be recalculated
        end
        GroupWeight(partitionable==0,:) = 0; % remove the non-partitionable regions
        GroupWeight(b,2) = 0;
        re_cal_id = find(GroupWeight(1:NumRegion,2)==1);
        if ~isempty(re_cal_id)
            re_cal_id = [b,re_cal_id'];
            %% from here only partitionable regions are considered
            K = length(re_cal_id); % number of groups
            n1_mod = n1(re_cal_id);
            cv_hat = sigma_hat(re_cal_id)./(mu_hat(re_cal_id)-y);
            if sigma_hat(b)<zero_tolerance
                cv_hat(1) = -1/norminv(Percentage,0,1); % if the best group has 0 variance
            end
            % the ration between group k and group b
            Wk = 1./(1+1./(cv_hat.^2)-1./n1_mod);
            Ckb = Wk./Wk(1);
            P_kb = fcdf(Ckb, n1_mod-1, ones(1,K).*n1_mod(1)-1); % the probability that cv_k>cv_b
            P_hat = P_kb./(1-P_kb); % Ni / Nb
            if partitionable(b)==0
                P_hat(1) = 0;
            end
            GroupWeight(re_cal_id,1) = P_hat';            
        end
        %% new budget allcation
        Ntot = sum(n1)+N_new;
        n = GroupWeight(1:NumRegion,1)./sum(GroupWeight(1:NumRegion,1)).*Ntot; %total budgets allocation
        n = n';
        n2 = zeros(1,NumRegion);
        % additional budgets allocation
        N_temp1 = N_new;
        [n_delta, k] = max(n-n1-n2);
        while N_temp1>0
            if n_delta<1
                add = 1;
            else
                add = max(0, min(N_temp1,round(n_delta)));
            end
            n2(k) = n2(k)+add;
            N_temp1 = N_temp1-add;
            [n_delta, k] = max(n-n1-n2);
        end 
        GroupWeight(:,2) = 0;
    end
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    function [Optima] = extract()
    %Nitching: obtain the final optima set
        %% ranking the solutions
        [~, Ranking] = sort(SampleSet(1:N0,2));
        SampleSet_temp = SampleSet(Ranking,:); % sort the samples from bad to good
        OptimaSet = OptCandidate(Ranking,:);
        %% clearing
        sampleid = 1;
        while 1
            Distance = (sqrt(sum((repmat(SampleSet_temp(sampleid,3:end),[N0,1])-SampleSet_temp(:,3:end)).^2,2)))';
            neighbor = find(Distance<=radius);
            [~,best_neighbor] = min(SampleSet_temp(neighbor,2));
            OptimaSet(neighbor) = 0; % delete all neighbor
            OptimaSet(neighbor(best_neighbor)) = 1; % save the best neighbor to the optima set
            if sampleid == (neighbor(best_neighbor))
                next = find(OptimaSet((sampleid+1):end),1);
                if isempty(next)
                    break
                else
                    sampleid = sampleid + next;
                end
            else
                sampleid = (neighbor(best_neighbor));
            end
        end
        Optima = SampleSet_temp(OptimaSet==1,:);
        OptCandidate = zeros(N,1);
        OptCandidate(Optima(:,1))=1;
    end
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
end

